from chromadb import PersistentClient
from sentence_transformers import SentenceTransformer
from rank_bm25 import BM25Okapi
import jieba
import numpy as np

# 初始化ChromaDB、向量模型、BM25
chroma_client = PersistentClient(path="./chroma_db")
collection = chroma_client.get_or_create_collection(name="rag_knowledge_base")
embedding_model = SentenceTransformer("all-MiniLM-L6-v2")  # 轻量级向量模型

# 知识库数据（示例，实际需从文件/数据库导入）
knowledge_base = [
    "产品A的定价为999元，支持30天无理由退货",
    "产品A的核心功能包括智能识别、自动同步、多设备兼容",
    "产品B的定价为1599元，支持1年质保"
]
# 为BM25构建分词后的语料
tokenized_corpus = [list(jieba.cut(doc)) for doc in knowledge_base]
bm25 = BM25Okapi(tokenized_corpus)


def hybrid_search(query: str, top_k: int = 3):
    # 1. 向量检索
    query_embedding = embedding_model.encode(query)
    vector_results = collection.query(
        query_embeddings=[query_embedding],
        n_results=top_k,
        include=["documents", "distances"]
    )
    # 提取向量检索结果（文档+相似度得分，距离越小相似度越高）
    vector_docs = vector_results["documents"][0]
    vector_scores = [1 - d for d in vector_results["distances"][0]]  # 距离转相似度

    # 2. BM25关键词检索
    query_tokens = list(jieba.cut(query))
    bm25_scores = bm25.get_scores(query_tokens)
    # 提取BM25结果（文档+得分）
    bm25_ranks = np.argsort(bm25_scores)[::-1][:top_k]  # 按得分降序取Top-K
    bm25_docs = [knowledge_base[i] for i in bm25_ranks]
    bm25_scores = [bm25_scores[i] for i in bm25_ranks]

    # 3. 结果合并与加权排序（向量得分占60%，BM25得分占40%）
    all_docs = list(set(vector_docs + bm25_docs))  # 去重
    final_scores = []
    for doc in all_docs:
        # 获取文档在两种检索中的得分（无则为0）
        v_score = vector_scores[vector_docs.index(doc)] if doc in vector_docs else 0
        b_score = bm25_scores[bm25_docs.index(doc)] if doc in bm25_docs else 0
        # 加权得分（需归一化，此处简化处理）
        weighted_score = v_score * 0.6 + b_score * 0.4
        final_scores.append((doc, weighted_score))

    # 按加权得分降序返回Top-K
    final_results = sorted(final_scores, key=lambda x: x[1], reverse=True)[:top_k]
    return [doc for doc, score in final_results]


# 示例：检索“产品A的定价”
if __name__ == "__main__":
    query = "产品A的定价是多少"
    results = hybrid_search(query)
    print("检索结果：", results)
    # 输出：检索结果：["产品A的定价为999元，支持30天无理由退货", "产品A的核心功能包括智能识别、自动同步、多设备兼容"]